import argparse
import os
import sys
from config import get_config
from logger import setup_logger

logger = setup_logger('run_server')


def start_http_server():
    """启动HTTP服务（使用FastAPI和Uvicorn）"""
    try:
        import uvicorn
        from api import app
        from model import sam_model
        
        config = get_config()
        
        logger.info(f"启动HTTP服务，监听地址: {config.HOST}:{config.PORT}, 设备: {config.DEVICE}")
        
        # 预加载模型
        logger.info("预加载SAM模型...")
        sam_model.load_model()
        
        # 启动服务
        uvicorn.run(
            "api:app",
            host=config.HOST,
            port=config.PORT,
            reload=False,
            workers=1,
            log_level=config.LOG_LEVEL.lower()
        )
    except ImportError as e:
        logger.error(f"导入依赖失败: {e}")
        logger.error("请先安装依赖: pip install -r requirements.txt")
        sys.exit(1)
    except Exception as e:
        logger.error(f"启动HTTP服务失败: {e}")
        sys.exit(1)

def start_ray_serve():
    """启动Ray Serve服务"""
    try:
        from ray_serve import start_ray_serve
        
        start_ray_serve()
    except ImportError as e:
        logger.error(f"导入依赖失败: {e}")
        logger.error("请先安装依赖: pip install -r requirements.txt")
        sys.exit(1)
    except Exception as e:
        logger.error(f"启动Ray Serve服务失败: {e}")
        sys.exit(1)

def download_model(model_type='vit_h', save_dir='./models'):
    """下载SAM模型（如果需要）"""
    try:
        import requests
        import os
        import tqdm
        
        # 模型下载链接
        model_urls = {
            'vit_h': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
            'vit_l': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
            'vit_b': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
            'mobile_sam': 'https://github.com/ChaoningZhang/MobileSAM/releases/download/v1.0/mobile_sam.pth'
        }
        
        if model_type not in model_urls:
            logger.error(f"不支持的模型类型: {model_type}")
            logger.error(f"支持的模型类型: {list(model_urls.keys())}")
            sys.exit(1)
        
        # 确保保存目录存在
        os.makedirs(save_dir, exist_ok=True)
        
        # 获取模型文件名和保存路径
        model_filename = model_urls[model_type].split('/')[-1]
        save_path = os.path.join(save_dir, model_filename)
        
        # 检查模型是否已存在
        if os.path.exists(save_path):
            logger.info(f"模型已存在: {save_path}")
            return save_path
        
        logger.info(f"下载模型: {model_type} ({model_urls[model_type]})")
        logger.info(f"保存路径: {save_path}")
        
        # 下载模型
        response = requests.get(model_urls[model_type], stream=True)
        response.raise_for_status()
        
        # 获取文件大小
        total_size = int(response.headers.get('content-length', 0))
        
        # 使用tqdm显示下载进度
        with open(save_path, 'wb') as f:
            with tqdm.tqdm(total=total_size, unit='B', unit_scale=True, desc=model_filename) as pbar:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        pbar.update(len(chunk))
        
        logger.info(f"模型下载完成: {save_path}")
        return save_path
    except Exception as e:
        logger.error(f"下载模型失败: {e}")
        sys.exit(1)

def main():
    """主函数"""
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='SAM Embedding Service')
    parser.add_argument('--mode', type=str, default='http', choices=['http', 'ray'],
                        help='服务运行模式: http (使用FastAPI) 或 ray (使用Ray Serve)')
    parser.add_argument('--download-model', type=str, choices=['vit_h', 'vit_l', 'vit_b', 'mobile_sam'],
                        help='下载指定的SAM模型')
    parser.add_argument('--model-dir', type=str, default='./models',
                        help='模型保存目录')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 如果指定了下载模型
    if args.download_model:
        download_model(args.download_model, args.model_dir)
        return
    
    # 根据模式启动服务
    if args.mode == 'http':
        start_http_server()
    elif args.mode == 'ray':
        start_ray_serve()

if __name__ == "__main__":
    main()